package jwt

import (
	"crypto/rsa"
	"errors"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
	"time"

	"gitcode.com/xiayang/faygo/v2/internal/context"
	"gopkg.in/dgrijalva/jwt-go.v3"
)

// HandlerFunc is a function type that handles HTTP requests with error return
type HandlerFunc func(*context.Context) error

// Map represents a generic map for JSON responses with type safety
type Map map[string]interface{}

// Claims represents JWT claims with type safety
type Claims map[string]interface{}

// UserData represents authenticated user data interface
type UserData interface{}

// TokenString represents a JWT token string with validation
type TokenString string

// ValidateToken validates token format
func (ts TokenString) IsValid() bool {
	return len(string(ts)) > 0 && strings.Count(string(ts), ".") == 2
}

// String returns the underlying string
func (ts TokenString) String() string {
	return string(ts)
}

// SigningMethod represents supported signing algorithms
type SigningMethod string

const (
	// HMAC signing methods
	HS256 SigningMethod = "HS256"
	HS384 SigningMethod = "HS384"
	HS512 SigningMethod = "HS512"
	// RSA signing methods
	RS256 SigningMethod = "RS256"
	RS384 SigningMethod = "RS384"
	RS512 SigningMethod = "RS512"
)

// IsValid checks if the signing method is supported
func (sm SigningMethod) IsValid() bool {
	switch sm {
	case HS256, HS384, HS512, RS256, RS384, RS512:
		return true
	default:
		return false
	}
}

// IsAsymmetric returns true if the method uses public/private key pairs
func (sm SigningMethod) IsAsymmetric() bool {
	switch sm {
	case RS256, RS384, RS512:
		return true
	default:
		return false
	}
}

// String returns the string representation
func (sm SigningMethod) String() string {
	return string(sm)
}

// TokenSource represents where to extract token from request
type TokenSource string

const (
	TokenSourceHeader TokenSource = "header"
	TokenSourceQuery  TokenSource = "query"
	TokenSourceCookie TokenSource = "cookie"
)

// AuthenticatorFunc defines the authentication callback function type
type AuthenticatorFunc func(userID, password string, c *context.Context) (UserData, bool)

// AuthorizatorFunc defines the authorization callback function type
type AuthorizatorFunc func(data UserData, c *context.Context) bool

// PayloadFunc defines the payload enhancement function type
type PayloadFunc func(data UserData) Claims

// UnauthorizedFunc defines the unauthorized response function type
type UnauthorizedFunc func(*context.Context, int, string)

// LoginResponseFunc defines the login response function type
type LoginResponseFunc func(*context.Context, int, TokenString, time.Time) error

// RefreshResponseFunc defines the refresh response function type
type RefreshResponseFunc func(*context.Context, int, TokenString, time.Time) error

// IdentityHandlerFunc defines the identity extraction function type
type IdentityHandlerFunc func(jwt.MapClaims) UserData

// HTTPStatusMessageFunc defines the error message function type
type HTTPStatusMessageFunc func(error, *context.Context) string

// TimeFunc defines the time function type for testing purposes
type TimeFunc func() time.Time

// Config represents JWT middleware configuration with strong typing
type Config struct {
	// Realm name to display to the user. Required.
	Realm string `json:"realm" validate:"required"`

	// Signing algorithm with type safety
	SigningAlgorithm SigningMethod `json:"signing_algorithm"`

	// Secret key used for signing. Required for HMAC algorithms.
	Key []byte `json:"-"` // Hide from JSON serialization for security

	// Duration that a jwt token is valid. Optional, defaults to one hour.
	Timeout time.Duration `json:"timeout"`

	// Maximum refresh duration
	MaxRefresh time.Duration `json:"max_refresh"`

	// Token lookup configuration
	TokenLookup string `json:"token_lookup"`

	// Token head name (Bearer, JWT, etc.)
	TokenHeadName string `json:"token_head_name"`

	// Key files for asymmetric algorithms
	PrivKeyFile string `json:"priv_key_file"`
	PubKeyFile  string `json:"pub_key_file"`

	// Parsed keys for asymmetric algorithms
	privKey *rsa.PrivateKey
	pubKey  *rsa.PublicKey
}

// Validate validates the configuration
func (c *Config) Validate() error {
	if c.Realm == "" {
		return ErrMissingRealm
	}

	if !c.SigningAlgorithm.IsValid() {
		return ErrInvalidSigningAlgorithm
	}

	if c.SigningAlgorithm.IsAsymmetric() {
		if c.PrivKeyFile == "" || c.PubKeyFile == "" {
			return errors.New("private and public key files are required for asymmetric algorithms")
		}
	} else {
		if len(c.Key) == 0 {
			return ErrMissingSecretKey
		}
	}

	return nil
}

// SetDefaults sets default values for the configuration
func (c *Config) SetDefaults() {
	if c.SigningAlgorithm == "" {
		c.SigningAlgorithm = HS256
	}

	if c.Timeout == 0 {
		c.Timeout = time.Hour
	}

	if c.TokenLookup == "" {
		c.TokenLookup = "header:Authorization"
	}

	if c.TokenHeadName == "" {
		c.TokenHeadName = "Bearer"
	}
}

// Callbacks represents JWT middleware callbacks with strong typing
type Callbacks struct {
	// Authentication callback
	Authenticator AuthenticatorFunc

	// Authorization callback
	Authorizator AuthorizatorFunc

	// Payload enhancement callback
	PayloadFunc PayloadFunc

	// Error response callbacks
	Unauthorized UnauthorizedFunc

	// Success response callbacks
	LoginResponse   LoginResponseFunc
	RefreshResponse RefreshResponseFunc

	// Identity extraction callback
	IdentityHandler IdentityHandlerFunc

	// Time function for testing
	TimeFunc TimeFunc

	// HTTP status message callback
	HTTPStatusMessageFunc HTTPStatusMessageFunc
}

// SetDefaults sets default callback implementations
func (cb *Callbacks) SetDefaults() {
	if cb.Authorizator == nil {
		cb.Authorizator = func(data UserData, c *context.Context) bool {
			return true
		}
	}

	if cb.Unauthorized == nil {
		cb.Unauthorized = func(c *context.Context, code int, message string) {
			c.JSON(code, Map{
				"code":    code,
				"message": message,
			})
		}
	}

	if cb.LoginResponse == nil {
		cb.LoginResponse = func(c *context.Context, code int, token TokenString, expire time.Time) error {
			return c.JSON(http.StatusOK, Map{
				"code":   http.StatusOK,
				"token":  token.String(),
				"expire": expire.Format(time.RFC3339),
			})
		}
	}

	if cb.RefreshResponse == nil {
		cb.RefreshResponse = func(c *context.Context, code int, token TokenString, expire time.Time) error {
			return c.JSON(http.StatusOK, Map{
				"code":   http.StatusOK,
				"token":  token.String(),
				"expire": expire.Format(time.RFC3339),
			})
		}
	}

	if cb.IdentityHandler == nil {
		cb.IdentityHandler = func(claims jwt.MapClaims) UserData {
			return claims["id"]
		}
	}

	if cb.TimeFunc == nil {
		cb.TimeFunc = time.Now
	}

	if cb.HTTPStatusMessageFunc == nil {
		cb.HTTPStatusMessageFunc = func(e error, c *context.Context) string {
			return e.Error()
		}
	}
}

// Validate validates the callbacks
func (cb *Callbacks) Validate() error {
	if cb.Authenticator == nil {
		return ErrMissingAuthenticatorFunc
	}
	return nil
}

// FaygoJWTMiddleware provides a Json-Web-Token authentication implementation with strong typing.
// On failure, a 401 HTTP response is returned. On success, the wrapped middleware is called,
// and the userID is made available as c.Data("userID").
// Users can get a token by posting a json request to LoginHandler. The token then needs to be passed in
// the Authentication header. Example: Authorization:Bearer XXX_TOKEN_XXX
type FaygoJWTMiddleware struct {
	// Configuration with validation
	Config *Config `json:"config"`

	// Callbacks with strong typing
	Callbacks *Callbacks `json:"-"` // Hide from JSON for security

	// Internal state
	initialized bool
}

var (
	// ErrMissingRealm indicates Realm name is required
	ErrMissingRealm = errors.New("realm is missing")

	// ErrMissingSecretKey indicates Secret key is required
	ErrMissingSecretKey = errors.New("secret key is required")

	// ErrForbidden when HTTP status 403 is given
	ErrForbidden = errors.New("you don't have permission to access this resource")

	// ErrMissingAuthenticatorFunc indicates Authenticator is required
	ErrMissingAuthenticatorFunc = errors.New("FaygoJWTMiddleware.Authenticator func is undefined")

	// ErrMissingLoginValues indicates a user tried to authenticate without username or password
	ErrMissingLoginValues = errors.New("missing Usercode or Password")

	// ErrFailedAuthentication indicates authentication failed, could be faulty username or password
	ErrFailedAuthentication = errors.New("incorrect Usercode or Password")

	// ErrFailedTokenCreation indicates JWT Token failed to create, reason unknown
	ErrFailedTokenCreation = errors.New("failed to create JWT Token")

	// ErrExpiredToken indicates JWT token has expired. Can't refresh.
	ErrExpiredToken = errors.New("token is expired")

	// ErrEmptyAuthHeader can be thrown if authing with a HTTP header, the Auth header needs to be set
	ErrEmptyAuthHeader = errors.New("auth header is empty")

	// ErrInvalidAuthHeader indicates auth header is invalid, could for example have the wrong Realm name
	ErrInvalidAuthHeader = errors.New("auth header is invalid")

	// ErrEmptyQueryToken can be thrown if authing with URL Query, the query token variable is empty
	ErrEmptyQueryToken = errors.New("query token is empty")

	// ErrEmptyCookieToken can be thrown if authing with a cookie, the token cokie is empty
	ErrEmptyCookieToken = errors.New("cookie token is empty")

	// ErrInvalidSigningAlgorithm indicates signing algorithm is invalid, needs to be HS256, HS384, HS512, RS256, RS384 or RS512
	ErrInvalidSigningAlgorithm = errors.New("invalid signing algorithm")

	// ErrNoPrivKeyFile indicates that the given private key is unreadable
	ErrNoPrivKeyFile = errors.New("private key file unreadable")

	// ErrNoPubKeyFile indicates that the given public key is unreadable
	ErrNoPubKeyFile = errors.New("public key file unreadable")

	// ErrInvalidPrivKey indicates that the given private key is invalid
	ErrInvalidPrivKey = errors.New("private key invalid")

	// ErrInvalidPubKey indicates the the given public key is invalid
	ErrInvalidPubKey = errors.New("public key invalid")
)

// LoginRequest represents the login request structure with validation
type LoginRequest struct {
	Usercode string `form:"usercode" json:"usercode" binding:"required" validate:"required,min=3,max=50"`
	Password string `form:"password" json:"password" binding:"required" validate:"required,min=6"`
}

// Validate validates the login request
func (lr *LoginRequest) Validate() error {
	if strings.TrimSpace(lr.Usercode) == "" {
		return errors.New("usercode is required")
	}
	if strings.TrimSpace(lr.Password) == "" {
		return errors.New("password is required")
	}
	if len(lr.Usercode) < 3 || len(lr.Usercode) > 50 {
		return errors.New("usercode must be between 3 and 50 characters")
	}
	if len(lr.Password) < 6 {
		return errors.New("password must be at least 6 characters")
	}
	return nil
}

// TokenResponse represents the token response structure
type TokenResponse struct {
	Code   int    `json:"code"`
	Token  string `json:"token"`
	Expire string `json:"expire"`
}

// ErrorResponse represents the error response structure
type ErrorResponse struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
	Error   string `json:"error,omitempty"`
}

// NewJWTMiddleware creates a new JWT middleware with proper configuration
func NewJWTMiddleware(config *Config, callbacks *Callbacks) *FaygoJWTMiddleware {
	if config == nil {
		config = &Config{}
	}
	if callbacks == nil {
		callbacks = &Callbacks{}
	}

	config.SetDefaults()
	callbacks.SetDefaults()

	return &FaygoJWTMiddleware{
		Config:    config,
		Callbacks: callbacks,
	}
}

// NewJWTMiddlewareFromLegacy creates middleware from legacy parameters for backward compatibility
func NewJWTMiddlewareFromLegacy(realm, signingAlgorithm string, key []byte, timeout time.Duration) *FaygoJWTMiddleware {
	config := &Config{
		Realm:            realm,
		SigningAlgorithm: SigningMethod(signingAlgorithm),
		Key:              key,
		Timeout:          timeout,
	}

	return NewJWTMiddleware(config, nil)
}

func (mw *FaygoJWTMiddleware) readKeys() error {
	err := mw.privateKey()
	if err != nil {
		return err
	}
	err = mw.publicKey()
	if err != nil {
		return err
	}
	return nil
}

func (mw *FaygoJWTMiddleware) privateKey() error {
	keyData, err := ioutil.ReadFile(mw.Config.PrivKeyFile)
	if err != nil {
		return ErrNoPrivKeyFile
	}
	key, err := jwt.ParseRSAPrivateKeyFromPEM(keyData)
	if err != nil {
		return ErrInvalidPrivKey
	}
	mw.Config.privKey = key
	return nil
}

func (mw *FaygoJWTMiddleware) publicKey() error {
	keyData, err := ioutil.ReadFile(mw.Config.PubKeyFile)
	if err != nil {
		return ErrNoPubKeyFile
	}
	key, err := jwt.ParseRSAPublicKeyFromPEM(keyData)
	if err != nil {
		return ErrInvalidPubKey
	}
	mw.Config.pubKey = key
	return nil
}

func (mw *FaygoJWTMiddleware) usingPublicKeyAlgo() bool {
	return mw.Config.SigningAlgorithm.IsAsymmetric()
}

// MiddlewareInit initialize jwt configs with enhanced validation.
func (mw *FaygoJWTMiddleware) MiddlewareInit() error {
	if mw.Config == nil {
		mw.Config = &Config{}
	}
	if mw.Callbacks == nil {
		mw.Callbacks = &Callbacks{}
	}

	// Set defaults
	mw.Config.SetDefaults()
	mw.Callbacks.SetDefaults()

	// Validate configuration
	if err := mw.Config.Validate(); err != nil {
		return err
	}

	// Validate callbacks
	if err := mw.Callbacks.Validate(); err != nil {
		return err
	}

	// Initialize keys for asymmetric algorithms
	if mw.usingPublicKeyAlgo() {
		if err := mw.readKeys(); err != nil {
			return err
		}
	}

	mw.initialized = true
	return nil
}

// MiddlewareFunc makes FaygoJWTMiddleware implement the Middleware interface.
func (mw *FaygoJWTMiddleware) MiddlewareFunc() HandlerFunc {
	return func(c *context.Context) error {
		return mw.middlewareImpl(c)
	}
}

func (mw *FaygoJWTMiddleware) middlewareImpl(c *context.Context) error {
	if !mw.initialized {
		if err := mw.MiddlewareInit(); err != nil {
			return err
		}
	}

	token, err := mw.parseToken(c)
	if err != nil {
		return mw.unauthorized(c, http.StatusUnauthorized, mw.Callbacks.HTTPStatusMessageFunc(err, c))
	}

	claims := token.Claims.(jwt.MapClaims)
	id := mw.Callbacks.IdentityHandler(claims)
	c.SetData("JWT_PAYLOAD", claims)
	c.SetData("userID", id)

	if !mw.Callbacks.Authorizator(id, c) {
		return mw.unauthorized(c, http.StatusForbidden, mw.Callbacks.HTTPStatusMessageFunc(ErrForbidden, c))
	}

	return nil
}

// LoginHandler can be used by clients to get a jwt token with enhanced validation.
// Payload needs to be json in the form of {"usercode": "USERNAME", "password": "PASSWORD"}.
// Reply will be of the form {"token": "TOKEN"}.
func (mw *FaygoJWTMiddleware) LoginHandler(c *context.Context) error {
	if !mw.initialized {
		if err := mw.MiddlewareInit(); err != nil {
			return err
		}
	}

	var loginVals LoginRequest
	if err := c.BindJSON(&loginVals); err != nil {
		return mw.unauthorized(c, http.StatusBadRequest, mw.Callbacks.HTTPStatusMessageFunc(ErrMissingLoginValues, c))
	}

	// Validate login request
	if err := loginVals.Validate(); err != nil {
		return mw.unauthorized(c, http.StatusBadRequest, mw.Callbacks.HTTPStatusMessageFunc(err, c))
	}

	if mw.Callbacks.Authenticator == nil {
		return mw.unauthorized(c, http.StatusInternalServerError, mw.Callbacks.HTTPStatusMessageFunc(ErrMissingAuthenticatorFunc, c))
	}

	data, ok := mw.Callbacks.Authenticator(loginVals.Usercode, loginVals.Password, c)
	if !ok {
		return mw.unauthorized(c, http.StatusUnauthorized, mw.Callbacks.HTTPStatusMessageFunc(ErrFailedAuthentication, c))
	}

	// Create the token
	token := jwt.New(jwt.GetSigningMethod(mw.Config.SigningAlgorithm.String()))
	claims := token.Claims.(jwt.MapClaims)

	if mw.Callbacks.PayloadFunc != nil {
		for key, value := range mw.Callbacks.PayloadFunc(data) {
			claims[key] = value
		}
	}
	if claims["id"] == nil {
		claims["id"] = loginVals.Usercode
	}
	expire := mw.Callbacks.TimeFunc().Add(mw.Config.Timeout)
	claims["exp"] = expire.Unix()
	claims["orig_iat"] = mw.Callbacks.TimeFunc().Unix()
	tokenString, err := mw.signedString(token)

	if err != nil {
		return mw.unauthorized(c, http.StatusUnauthorized, mw.Callbacks.HTTPStatusMessageFunc(ErrFailedTokenCreation, c))
	}

	return mw.Callbacks.LoginResponse(c, http.StatusOK, TokenString(tokenString), expire)
}

func (mw *FaygoJWTMiddleware) signedString(token *jwt.Token) (string, error) {
	var tokenString string
	var err error
	if mw.usingPublicKeyAlgo() {
		tokenString, err = token.SignedString(mw.Config.privKey)
	} else {
		tokenString, err = token.SignedString(mw.Config.Key)
	}
	return tokenString, err
}

// RefreshHandler can be used to refresh a token with enhanced validation. The token still needs to be valid on refresh.
// Shall be put under an endpoint that is using the FaygoJWTMiddleware.
// Reply will be of the form {"token": "TOKEN"}.
func (mw *FaygoJWTMiddleware) RefreshHandler(c *context.Context) error {
	if !mw.initialized {
		if err := mw.MiddlewareInit(); err != nil {
			return err
		}
	}

	token, err := mw.parseToken(c)
	// 如果有错误并不是过期则触发错误并返回
	if (err != nil) && (err.Error() != "Token is expired") {
		log.Printf("JWT refresh error: %v", err)
		return mw.unauthorized(c, http.StatusUnauthorized, mw.Callbacks.HTTPStatusMessageFunc(err, c))
	}
	claims := token.Claims.(jwt.MapClaims)

	origIat := int64(claims["orig_iat"].(float64))
	// 如果超过刷新期则触发过期错误并返回
	if origIat < mw.Callbacks.TimeFunc().Add(-mw.Config.MaxRefresh).Unix() {
		return mw.unauthorized(c, http.StatusUnauthorized, mw.Callbacks.HTTPStatusMessageFunc(ErrExpiredToken, c))
	}

	// Create the token
	newToken := jwt.New(jwt.GetSigningMethod(mw.Config.SigningAlgorithm.String()))
	newClaims := newToken.Claims.(jwt.MapClaims)

	for key := range claims {
		newClaims[key] = claims[key]
	}

	expire := mw.Callbacks.TimeFunc().Add(mw.Config.Timeout)
	newClaims["id"] = claims["id"]
	newClaims["exp"] = expire.Unix()
	newClaims["orig_iat"] = mw.Callbacks.TimeFunc().Unix()
	tokenString, err := mw.signedString(newToken)

	if err != nil {
		return mw.unauthorized(c, http.StatusUnauthorized, mw.Callbacks.HTTPStatusMessageFunc(ErrFailedTokenCreation, c))
	}

	return mw.Callbacks.RefreshResponse(c, http.StatusOK, TokenString(tokenString), expire)
}

// ExtractClaims help to extract the JWT claims
func ExtractClaims(c *context.Context) jwt.MapClaims {
	if exists := c.HasData("JWT_PAYLOAD"); !exists {
		return make(jwt.MapClaims)
	}

	jwtClaims := c.Data("JWT_PAYLOAD")
	return jwtClaims.(jwt.MapClaims)
}

// TokenGenerator method that clients can use to get a jwt token with enhanced validation.
func (mw *FaygoJWTMiddleware) TokenGenerator(userID string) (TokenString, time.Time, error) {
	if !mw.initialized {
		if err := mw.MiddlewareInit(); err != nil {
			return "", time.Time{}, err
		}
	}

	token := jwt.New(jwt.GetSigningMethod(mw.Config.SigningAlgorithm.String()))
	claims := token.Claims.(jwt.MapClaims)

	if mw.Callbacks.PayloadFunc != nil {
		for key, value := range mw.Callbacks.PayloadFunc(userID) {
			claims[key] = value
		}
	}

	expire := mw.Callbacks.TimeFunc().UTC().Add(mw.Config.Timeout)
	claims["id"] = userID
	claims["exp"] = expire.Unix()
	claims["orig_iat"] = mw.Callbacks.TimeFunc().Unix()
	tokenString, err := mw.signedString(token)
	if err != nil {
		return "", time.Time{}, err
	}

	return TokenString(tokenString), expire, nil
}

func (mw *FaygoJWTMiddleware) jwtFromHeader(c *context.Context, key string) (string, error) {
	authHeader := c.HeaderParam(key)

	if authHeader == "" {
		return "", ErrEmptyAuthHeader
	}

	parts := strings.SplitN(authHeader, " ", 2)
	if !(len(parts) == 2 && parts[0] == mw.Config.TokenHeadName) {
		return "", ErrInvalidAuthHeader
	}

	return parts[1], nil
}

func (mw *FaygoJWTMiddleware) jwtFromQuery(c *context.Context, key string) (string, error) {
	token := c.QueryParam(key)

	if token == "" {
		return "", ErrEmptyQueryToken
	}

	return token, nil
}

func (mw *FaygoJWTMiddleware) jwtFromCookie(c *context.Context, key string) (string, error) {
	cookie := c.CookieParam(key)

	if cookie == "" {
		return "", ErrEmptyCookieToken
	}

	return cookie, nil
}

func (mw *FaygoJWTMiddleware) parseToken(c *context.Context) (*jwt.Token, error) {
	var token string
	var err error

	parts := strings.Split(mw.Config.TokenLookup, ":")
	switch parts[0] {
	case "header":
		token, err = mw.jwtFromHeader(c, parts[1])
	case "query":
		token, err = mw.jwtFromQuery(c, parts[1])
	case "cookie":
		token, err = mw.jwtFromCookie(c, parts[1])
	}

	if err != nil {
		return nil, err
	}

	return jwt.Parse(token, func(token *jwt.Token) (interface{}, error) {
		if jwt.GetSigningMethod(mw.Config.SigningAlgorithm.String()) != token.Method {
			return nil, ErrInvalidSigningAlgorithm
		}
		if mw.usingPublicKeyAlgo() {
			return mw.Config.pubKey, nil
		}
		return mw.Config.Key, nil
	})
}

func (mw *FaygoJWTMiddleware) unauthorized(c *context.Context, code int, message string) error {
	realm := mw.Config.Realm
	if realm == "" {
		realm = "faygo jwt"
	}
	c.SetHeader("WWW-Authenticate", "JWT realm="+realm)
	c.Stop()

	mw.Callbacks.Unauthorized(c, code, message)
	return nil
}
